# -*- coding: utf-8 -*-
"""
Created on 2019/11/22 上午10:22

@Project -> File: ode-neural-network -> nn.py

@Author: luolei

@Describe:
"""

import torch
from torch import nn
from torch.nn import init


class PartialDeriveNet(nn.Module):
	"""
	内层用于计算偏导的神经网络模型
	"""
	
	def __init__(self, input_size, hidden_sizes, output_size):
		super(PartialDeriveNet, self).__init__()
		self.input_size = input_size
		self.hidden_sizes = hidden_sizes
		self.output_size = output_size
		
		self.fc_0 = nn.Linear(self.input_size, self.hidden_sizes[0])
		self._init_layer(self.fc_0)
		
		self.fcs = []
		for i in range(len(hidden_sizes) - 1):
			fc_i = nn.Linear(self.hidden_sizes[i], self.hidden_sizes[i + 1])
			setattr(self, 'fc_{}'.format(i + 1), fc_i)
			self._init_layer(fc_i)
			self.fcs.append(fc_i)
		
		self.fc_out = nn.Linear(self.hidden_sizes[-1], self.output_size)
		self.prelu = nn.PReLU()  # todo: 这里面的 n_parameters 参数意义没有搞清楚
		self._init_layer(self.fc_out)
	
	def _init_layer(self, layer):
		init.normal_(layer.weight)  # 使用这种初始化方式能降低过拟合
		init.normal_(layer.bias)
	
	def forward(self, x):
		x = self.fc_0(x)
		x = torch.sigmoid(x)
		
		for i in range(len(self.fcs)):
			x = self.fcs[i](x)
			x = torch.sigmoid(x)
		
		x = self.fc_out(x)
		x = self.prelu(x)
		
		return x



